"""
语义检索模块

这个模块负责：
1. 语义检索功能
2. 检索结果排序和过滤
3. 检索策略优化

作者：RAG学习项目
"""

import logging
import numpy as np
from typing import List, Dict, Any, Optional, Tuple
from .config import config
from .embeddings import EmbeddingGenerator, VectorStore

# 配置日志
logging.basicConfig(level=getattr(logging, config.LOG_LEVEL))
logger = logging.getLogger(__name__)

class SemanticRetriever:
    """
    语义检索器类
    
    负责：
    - 执行语义相似度检索
    - 管理检索策略
    - 优化检索性能
    """
    
    def __init__(self, vector_store: VectorStore = None, embedding_generator: EmbeddingGenerator = None):
        """
        初始化语义检索器
        
        Args:
            vector_store (VectorStore, optional): 向量存储，如果不提供则创建新的
            embedding_generator (EmbeddingGenerator, optional): 向量生成器，如果不提供则创建新的
        """
        self.vector_store = vector_store or VectorStore()
        self.embedding_generator = embedding_generator or EmbeddingGenerator()
        
        logger.info("初始化语义检索器")
    
    def add_documents(self, documents: List[str], metadata: List[Dict[str, Any]] = None):
        """
        添加文档到检索库
        
        Args:
            documents (List[str]): 文档文本列表
            metadata (List[Dict[str, Any]], optional): 文档元数据列表
        """
        if not documents:
            logger.warning("文档列表为空")
            return
        
        logger.info(f"正在添加 {len(documents)} 个文档到检索库...")
        
        try:
            # 批量生成向量
            embeddings = self.embedding_generator.encode_texts(documents)
            
            # 添加到向量存储
            self.vector_store.add_documents(documents, embeddings, metadata)
            
            logger.info(f"✅ 成功添加 {len(documents)} 个文档")
            
        except Exception as e:
            logger.error(f"❌ 添加文档失败: {str(e)}")
            raise Exception(f"添加文档失败: {str(e)}")
    
    def search(self, query: str, top_k: int = None, similarity_threshold: float = None) -> List[Dict[str, Any]]:
        """
        执行语义检索
        
        Args:
            query (str): 查询文本
            top_k (int, optional): 返回最相关的K个文档，默认使用配置值
            similarity_threshold (float, optional): 相似度阈值，默认使用配置值
            
        Returns:
            List[Dict[str, Any]]: 检索结果列表
        """
        if not query or not query.strip():
            logger.warning("查询文本为空")
            return []
        
        # 使用配置的默认值
        top_k = top_k or config.TOP_K
        similarity_threshold = similarity_threshold or config.SIMILARITY_THRESHOLD
        
        logger.info(f"执行语义检索，查询: '{query[:50]}...', top_k: {top_k}, 阈值: {similarity_threshold}")
        
        try:
            # 将查询文本向量化
            query_embedding = self.embedding_generator.encode_text(query)
            
            # 在向量存储中搜索
            results = self.vector_store.search(
                query_embedding, 
                top_k=top_k, 
                similarity_threshold=similarity_threshold
            )
            
            logger.info(f"✅ 检索完成，找到 {len(results)} 个相关文档")
            
            # 添加查询信息到结果中
            for result in results:
                result['query'] = query
                result['search_params'] = {
                    'top_k': top_k,
                    'similarity_threshold': similarity_threshold
                }
            
            return results
            
        except Exception as e:
            logger.error(f"❌ 语义检索失败: {str(e)}")
            raise Exception(f"语义检索失败: {str(e)}")
    
    def get_document_count(self) -> int:
        """
        获取检索库中的文档数量
        
        Returns:
            int: 文档数量
        """
        return self.vector_store.get_document_count()
    
    def clear_documents(self):
        """清空检索库中的所有文档"""
        self.vector_store.clear()
        logger.info("检索库已清空")

class ResultRanker:
    """
    检索结果排序器类
    
    负责：
    - 对检索结果进行重新排序
    - 应用多种排序策略
    - 过滤和优化结果
    """
    
    def __init__(self):
        """初始化结果排序器"""
        logger.info("初始化检索结果排序器")
    
    def rank_by_similarity(self, results: List[Dict[str, Any]], reverse: bool = True) -> List[Dict[str, Any]]:
        """
        按相似度排序结果
        
        Args:
            results (List[Dict[str, Any]]): 检索结果列表
            reverse (bool): 是否降序排列，默认True（相似度高的在前）
            
        Returns:
            List[Dict[str, Any]]: 排序后的结果列表
        """
        if not results:
            return []
        
        logger.debug(f"按相似度排序 {len(results)} 个结果")
        
        try:
            sorted_results = sorted(
                results, 
                key=lambda x: x.get('similarity', 0), 
                reverse=reverse
            )
            
            # 更新排名
            for i, result in enumerate(sorted_results):
                result['rank'] = i + 1
            
            return sorted_results
            
        except Exception as e:
            logger.error(f"排序失败: {str(e)}")
            return results
    
    def filter_by_threshold(self, results: List[Dict[str, Any]], threshold: float) -> List[Dict[str, Any]]:
        """
        根据相似度阈值过滤结果
        
        Args:
            results (List[Dict[str, Any]]): 检索结果列表
            threshold (float): 相似度阈值
            
        Returns:
            List[Dict[str, Any]]: 过滤后的结果列表
        """
        if not results:
            return []
        
        logger.debug(f"使用阈值 {threshold} 过滤结果")
        
        filtered_results = [
            result for result in results 
            if result.get('similarity', 0) >= threshold
        ]
        
        logger.debug(f"过滤后保留 {len(filtered_results)}/{len(results)} 个结果")
        return filtered_results
    
    def filter_by_length(self, results: List[Dict[str, Any]], 
                        min_length: int = 10, max_length: int = 2000) -> List[Dict[str, Any]]:
        """
        根据文档长度过滤结果
        
        Args:
            results (List[Dict[str, Any]]): 检索结果列表
            min_length (int): 最小长度
            max_length (int): 最大长度
            
        Returns:
            List[Dict[str, Any]]: 过滤后的结果列表
        """
        if not results:
            return []
        
        logger.debug(f"按长度过滤结果，范围: {min_length}-{max_length} 字符")
        
        filtered_results = [
            result for result in results 
            if min_length <= len(result.get('document', '')) <= max_length
        ]
        
        logger.debug(f"长度过滤后保留 {len(filtered_results)}/{len(results)} 个结果")
        return filtered_results
    
    def deduplicate_results(self, results: List[Dict[str, Any]], 
                           similarity_threshold: float = 0.95) -> List[Dict[str, Any]]:
        """
        去除重复或高度相似的结果
        
        Args:
            results (List[Dict[str, Any]]): 检索结果列表
            similarity_threshold (float): 判断重复的相似度阈值
            
        Returns:
            List[Dict[str, Any]]: 去重后的结果列表
        """
        if not results or len(results) <= 1:
            return results
        
        logger.debug(f"对 {len(results)} 个结果进行去重")
        
        try:
            # 创建向量生成器用于计算文档间相似度
            embedding_generator = EmbeddingGenerator()
            
            unique_results = []
            
            for current_result in results:
                current_doc = current_result.get('document', '')
                is_duplicate = False
                
                # 与已选择的结果比较
                for unique_result in unique_results:
                    unique_doc = unique_result.get('document', '')
                    
                    # 计算文档间相似度
                    current_embedding = embedding_generator.encode_text(current_doc)
                    unique_embedding = embedding_generator.encode_text(unique_doc)
                    
                    doc_similarity = embedding_generator.calculate_similarity(
                        current_embedding, unique_embedding
                    )
                    
                    if doc_similarity >= similarity_threshold:
                        is_duplicate = True
                        break
                
                if not is_duplicate:
                    unique_results.append(current_result)
            
            logger.debug(f"去重后保留 {len(unique_results)}/{len(results)} 个结果")
            return unique_results
            
        except Exception as e:
            logger.error(f"去重失败: {str(e)}")
            return results
    
    def apply_comprehensive_ranking(self, results: List[Dict[str, Any]], 
                                  similarity_weight: float = 0.7,
                                  length_weight: float = 0.2,
                                  position_weight: float = 0.1) -> List[Dict[str, Any]]:
        """
        应用综合排序策略
        
        Args:
            results (List[Dict[str, Any]]): 检索结果列表
            similarity_weight (float): 相似度权重
            length_weight (float): 长度权重
            position_weight (float): 位置权重
            
        Returns:
            List[Dict[str, Any]]: 综合排序后的结果列表
        """
        if not results:
            return []
        
        logger.debug(f"应用综合排序策略，权重: 相似度{similarity_weight}, 长度{length_weight}, 位置{position_weight}")
        
        try:
            # 计算各项分数的正规化值
            similarities = [r.get('similarity', 0) for r in results]
            lengths = [len(r.get('document', '')) for r in results]
            
            max_similarity = max(similarities) if similarities else 1
            max_length = max(lengths) if lengths else 1
            min_length = min(lengths) if lengths else 0
            
            # 为每个结果计算综合分数
            for i, result in enumerate(results):
                similarity_score = result.get('similarity', 0) / max_similarity
                
                # 长度分数：中等长度的文档得分更高
                doc_length = len(result.get('document', ''))
                if max_length > min_length:
                    # 使用正态分布式的长度评分
                    optimal_length = (max_length + min_length) / 2
                    length_score = 1 - abs(doc_length - optimal_length) / (max_length - min_length)
                else:
                    length_score = 1
                
                # 位置分数：原始排名靠前的得分更高
                position_score = 1 - (i / len(results))
                
                # 综合分数
                comprehensive_score = (
                    similarity_weight * similarity_score +
                    length_weight * length_score +
                    position_weight * position_score
                )
                
                result['comprehensive_score'] = comprehensive_score
                result['similarity_score'] = similarity_score
                result['length_score'] = length_score
                result['position_score'] = position_score
            
            # 按综合分数排序
            sorted_results = sorted(
                results, 
                key=lambda x: x.get('comprehensive_score', 0), 
                reverse=True
            )
            
            # 更新排名
            for i, result in enumerate(sorted_results):
                result['rank'] = i + 1
            
            logger.debug("综合排序完成")
            return sorted_results
            
        except Exception as e:
            logger.error(f"综合排序失败: {str(e)}")
            return results

class RetrievalEngine:
    """
    检索引擎类
    
    整合语义检索器和结果排序器，提供完整的检索功能
    """
    
    def __init__(self):
        """初始化检索引擎"""
        self.retriever = SemanticRetriever()
        self.ranker = ResultRanker()
        
        logger.info("初始化检索引擎")
    
    def add_documents(self, documents: List[str], metadata: List[Dict[str, Any]] = None):
        """
        添加文档到检索引擎
        
        Args:
            documents (List[str]): 文档文本列表
            metadata (List[Dict[str, Any]], optional): 文档元数据列表
        """
        self.retriever.add_documents(documents, metadata)
    
    def search_and_rank(self, query: str, 
                       top_k: int = None,
                       similarity_threshold: float = None,
                       apply_filters: bool = True,
                       deduplicate: bool = True) -> List[Dict[str, Any]]:
        """
        执行检索并排序
        
        Args:
            query (str): 查询文本
            top_k (int, optional): 返回结果数量
            similarity_threshold (float, optional): 相似度阈值
            apply_filters (bool): 是否应用过滤器
            deduplicate (bool): 是否去重
            
        Returns:
            List[Dict[str, Any]]: 处理后的检索结果
        """
        logger.info(f"执行完整检索流程，查询: '{query[:50]}...'")
        
        try:
            # 1. 执行语义检索
            results = self.retriever.search(query, top_k, similarity_threshold)
            
            if not results:
                logger.info("未找到相关文档")
                return []
            
            # 2. 应用过滤器
            if apply_filters:
                # 长度过滤
                results = self.ranker.filter_by_length(results)
                
                # 相似度过滤
                threshold = similarity_threshold or config.SIMILARITY_THRESHOLD
                results = self.ranker.filter_by_threshold(results, threshold)
            
            # 3. 去重
            if deduplicate and len(results) > 1:
                results = self.ranker.deduplicate_results(results)
            
            # 4. 综合排序
            results = self.ranker.apply_comprehensive_ranking(results)
            
            logger.info(f"✅ 检索流程完成，最终返回 {len(results)} 个结果")
            return results
            
        except Exception as e:
            logger.error(f"❌ 检索流程失败: {str(e)}")
            raise Exception(f"检索流程失败: {str(e)}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取检索引擎统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        return {
            'document_count': self.retriever.get_document_count(),
            'embedding_dimension': self.retriever.embedding_generator.get_embedding_dimension(),
            'model_name': self.retriever.embedding_generator.model_name
        }
    
    def clear_all(self):
        """清空所有数据"""
        self.retriever.clear_documents()
        logger.info("检索引擎已清空")

def test_retrieval_system():
    """
    测试检索系统
    """
    print("=== 测试检索系统 ===")
    
    try:
        # 创建检索引擎
        engine = RetrievalEngine()
        
        # 准备测试文档
        test_documents = [
            "翁氏宗谱记录了家族的历史传承，包含了重要的族谱信息。",
            "家族中有许多杰出人物，为社会做出了重要贡献。",
            "族谱中详细记录了各代族人的姓名、生卒年月等信息。",
            "翁氏家族起源于古代，有着深厚的文化底蕴。",
            "现代族人分布在世界各地，但都保持着家族传统。"
        ]
        
        test_metadata = [
            {'source': 'page_1', 'type': 'general'},
            {'source': 'page_2', 'type': 'biography'},
            {'source': 'page_3', 'type': 'records'},
            {'source': 'page_4', 'type': 'history'},
            {'source': 'page_5', 'type': 'modern'}
        ]
        
        # 添加文档
        engine.add_documents(test_documents, test_metadata)
        print(f"✅ 添加了 {len(test_documents)} 个测试文档")
        
        # 测试检索
        query = "翁氏家族的历史"
        results = engine.search_and_rank(query, top_k=3)
        
        print(f"✅ 检索完成，找到 {len(results)} 个相关文档：")
        for result in results:
            print(f"   排名: {result['rank']}")
            print(f"   文档: {result['document'][:50]}...")
            print(f"   相似度: {result['similarity']:.4f}")
            print(f"   综合分数: {result.get('comprehensive_score', 0):.4f}")
            print(f"   元数据: {result['metadata']}")
            print()
        
        # 获取统计信息
        stats = engine.get_statistics()
        print(f"✅ 系统统计: {stats}")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    test_retrieval_system() 